"""
将数据集随机分成训练集、验证集
传入参数：
ratio: 训练样本比例
path: 数据路径
new_path: 保存路径
使用方法：
temp = Generate_Train_and_Val(path, new_path, ratio)
temp.splict_data()
"""

import random
import os
import shutil
from utils.ConfigReader import ConfigReader

def makeDir(path):
    """
    新建文件夹
    :param path: 路径
    :return: 0--创建成功，1--文件夹已存在，2--error
    """
    try:
        if not os.path.exists(path):
            os.makedirs(path)
            return 0
        else:
            return 1
    except Exception as e:
        print(str(e))
        return -2


def renameFiles(path):
    """
    按序重命名该路径下的数据集图片名
    :param path: 数据集路径
    :return: 是否重命名成功
    """
    if not os.path.isdir(path):
        print("{} is not a dir!".format(path))
        return False
    files_list = os.listdir(path)  # 返回path路径下的文件名列表(只是文件名，不是路径)
    assert os.path.isfile(os.path.join(path, files_list[0])), "path or files error!"
    files_list.sort(key=lambda x: int(x.split('.')[0]))  # 列表中的文件名排序

    # 将原路径的文件rename到该路径
    count = 0
    file_type = os.path.splitext(files_list[0])[1]  # 得到图片类型
    # 遍历文件并按序重命名
    for file in files_list:
        old_file = os.path.join(path, file)
        new_file = os.path.join(path, str(count) + file_type)
        # 若已存在该文件则跳过
        if os.path.exists(new_file):
            count += 1
            continue
        else:
            os.rename(old_file, new_file)
            count += 1
    return True


class Generate_Train_and_Val:
    def __init__(self, path, new_path, ratio):
        """
        训练集测试集打乱分类
        :param path: 图片路径
        :param new_path: 数据集路径
        :param ratio: 训练集比例
        """
        if not os.path.exists(new_path):
            makeDir(new_path)
        self.__path = path
        self.__new_path = new_path
        self.__ratio = ratio

        self.__train_sample_path = os.path.join(new_path, "train")
        self.__val_sample_path = os.path.join(new_path, "val")
        makeDir(self.__train_sample_path)
        makeDir(self.__val_sample_path)

    def splict_data(self):
        """
        数据集分离
        """
        class_names = os.listdir(self.__path)  # 类别文件夹
        delete_list = ['train', 'val', 'test', 'runs'] # 文件夹删除列表
        # 删除预先设置好的文件夹
        for i in delete_list:
            if i in class_names:
                class_names.remove(i)

        # 遍历类别名
        for name in class_names:
            print("\033[1mprocess class name: %s\033[0m" % name)
            tmp_class_name = os.path.join(self.__path, name)  # 某类别的原文件夹路径
            save_train_class_name = os.path.join(self.__train_sample_path, name)  # 某类别的训练集路径
            save_val_class_name = os.path.join(self.__val_sample_path, name)  # 某类别的验证集路径

            # 创建该类别的训练集文件夹
            makeDir(save_train_class_name)
            makeDir(save_val_class_name)

            # 数据集 ==> 训练集 验证集
            if os.path.isdir(tmp_class_name):
                image_names = os.listdir(tmp_class_name)  # 其中一个类别的所有图像
                total = len(image_names)  # 类别图片数

                # 1, 打乱当前类中所有图像
                random.shuffle(image_names)

                # 2, 从当前类中，取前面的图像作为train data
                train_temp = int(self.__ratio * total)  # 打乱后，取前面作为train_data
                for i in range(0, train_temp):
                    # 读取训练集图像
                    temp_img_name = os.path.join(tmp_class_name, image_names[i])
                    # 训练集图像保存路径
                    save_train_img_name = os.path.join(save_train_class_name, image_names[i])
                    shutil.copyfile(temp_img_name, save_train_img_name)

                # 3, 从当前类中，取后面的图像作为val data
                for i in range(train_temp, total):
                    # 数据集图像路径
                    val_img_name = os.path.join(tmp_class_name, image_names[i])
                    # 验证集图像保存路径
                    save_val_img_name = os.path.join(save_val_class_name, image_names[i])
                    shutil.copyfile(val_img_name, save_val_img_name)

            renameFiles(save_train_class_name)
            renameFiles(save_val_class_name)


def main():
    os.chdir('../')  # 工作路径切换为上一路径
    CR = ConfigReader('config.yaml')
    dataset_path = CR.getElement('dataset_path')

    generator = Generate_Train_and_Val(dataset_path, dataset_path, 0.7)
    generator.splict_data()


if __name__ == '__main__':
    main()
